# 导入必要的库和模块
import gradio as gr
import joblib
import numpy as np
from PIL import Image

# 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
knn_model = joblib.load("knn_model.pkl")  # 请确保模型文件名和路径正确

# 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(img):
    # 递归取出 image 数据
    while isinstance(img, dict):
        if "image" in img:
            img = img["image"]
        elif "layers" in img and isinstance(img["layers"], list) and len(img["layers"]) > 0:
            img = img["layers"][0]
        else:
            break
    # 转为PIL Image
    if isinstance(img, np.ndarray):
        # 如果是RGBA，转灰度
        if img.ndim == 3 and img.shape[2] == 4:
            img = Image.fromarray(img.astype('uint8'), mode="RGBA").convert("L")
        else:
            img = Image.fromarray(img.astype('uint8'), mode="L")
    # 关键：缩放到8x8
    img = img.resize((8, 8))
    img_arr = np.array(img)
    # 归一化到0~16
    img_arr = (img_arr / 255.0) * 16
    img_arr = img_arr.reshape(1, -1)
    # 预测
    pred = knn_model.predict(img_arr)
    return int(pred[0])

# 创建Gradio接口，这个接口将用于用户输入和显示预测结果
interface = gr.Interface(
    fn=predict_digit,
    inputs=gr.Sketchpad(canvas_size=(28, 28), label="请手写一个数字"),
    outputs=gr.Number(label="预测结果"),
    title="手写数字识别（KNN）",
    description="请在画板上手写一个数字，模型会预测你写的是几。"
)

# 启动Gradio接口，用户可以通过这个接口进行交互
if __name__ == "__main__":
    interface.launch()
